# !/usr/bin/env python
# -*- coding:utf-8 -*-
# @Project   :pyTools 
# @FileName  :socket.py
# @Time      :2023/12/1 13:00
# @Author    :zxc19
# @IDE       :PyCharm
# @DESC      :
import ipaddress
import logging
import os
import socket
import threading
from abc import abstractmethod

from greenlet import getcurrent
from locust import User, env


class SocketAbstractUser(User):
    abstract = True  # 定义为抽象用户，必须被继承
    client = None  # 定义客户端
    environment: env.Environment

    def __init__(self, environment):
        super().__init__(environment)
        self.process_id = None  # 进程ID
        self.thread_id = None  # 线程ID
        self.coroutine_id = None  # 协程ID
        # 获取当前用户的进程号
        self.process_id = os.getpid()
        # 获取当前用户的线程号
        self.thread_id = threading.current_thread().ident
        # 获取当前用户的协程号
        self.coroutine_id = id(getcurrent())

        logging.debug(
            f'用户初始化：{self.__class__.__name__}\t'
            f'Process ID: {self.process_id}\t'
            f'Thread ID: {self.thread_id}\t'
            f'Coroutine ID: {self.coroutine_id}\t'
        )

    def on_start(self):
        """
        用户启动
        :return:
        """
        logging.debug(
            f'用户启动：{self.__class__.__name__}\t'
            f'Process ID: {self.process_id}\t'
            f'Thread ID: {self.thread_id}\t'
            f'Coroutine ID: {self.coroutine_id}\t'
        )

    def on_stop(self):
        """
        用户停止
        :return:
        """

        logging.debug(
            f'用户停止：{self.__class__.__name__}\t'
            f'Process ID: {self.process_id}\t'
            f'Thread ID: {self.thread_id}\t'
            f'Coroutine ID: {self.coroutine_id}'
        )

    @staticmethod
    def checkip(address: str) -> str:
        """
        检查IP地址
        :param address: 地址
        :return: IPv4 / IPv6 / Unknown
        """
        try:
            ip = ipaddress.ip_address(address)
            if isinstance(ip, ipaddress.IPv4Address):
                return 'IPv4'
            elif isinstance(ip, ipaddress.IPv6Address):
                return 'IPv6'
        except ValueError:
            return 'Unknown'

    @abstractmethod
    def send(
            self,
            dst_ip: str, dst_port: int, data: str, local_ip: str = None, local_port: int = None,
            name: str = None, context: dict = None
    ):
        """
        发送消息
        :param dst_ip: 目的IP
        :param dst_port: 目的端口
        :param data: 数据(字符串)
        :param local_ip: 本地发送IP
        :param local_port: 本地发送端口
        :param name: 名称
        :param context: 上下文
        :return:
        """
        pass


class UDPSocketAbstractUser(SocketAbstractUser):
    abstract = True  # 定义为抽象用户，必须被继承

    def send(
            self, dst_ip: str, dst_port: int, data: str, local_ip: str = None, local_port: int = None,
            name: str = None, context: dict = None
    ):
        """
        发送消息UDP消息
        :param dst_ip: 目的IP
        :param dst_port: 目的端口
        :param data: 数据(字符串)
        :param local_ip: 本地发送IP
        :param local_port: 本地发送端口
        :param name: 名称
        :param context: 上下文
        :return:
        """
        if context is None:
            context = {}

        ip_type = self.checkip(dst_ip)
        if ip_type == 'IPv4':
            self.client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        elif ip_type == 'IPv6':
            self.client = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
        else:
            raise ValueError('目标IP不是一个有效的IP地址')
        if local_port:
            # 绑定发送的网卡和端口
            self.client.bind(('', local_port))
        self.environment.events.request.fire(
            request_type="UDP",
            name=name or dst_ip,
            response_time=None,
            response_length=len(data),
            exception=None,
            context={**self.context(), **context},
        )
        self.client.sendto(data.encode('utf-8'), (dst_ip, dst_port))


if __name__ == "__main__":
    run_code = 0
